#! /usr/bin/env python
"""
Utilities for generating Flatbuffers schemas for ROS message types. 

Schema formats documentation:
    http://wiki.ros.org/msg and 
    https://google.github.io/flatbuffers/flatbuffers_guide_writing_schema.html
"""

from __future__ import print_function, division
from argparse import ArgumentParser
import os
import re
import sys

# from rosbridge_library.internal.ros_loader import get_message_class
from roslib.message import get_message_class

# Flatbuffers types considered "scalar"
scalar_types = {
    "bool",
    "int8",
    "int16",
    "int32",
    "int64",
    "uint8",
    "uint16",
    "uint32",
    "uint64",
    "float32",
    "float64",
}

# types that are already defined in Flatbuffers
primitive_types = scalar_types | {
    "string",
}

# all types that can be used in Flatbuffers at startup
# all of the Flatbuffers primitives directly correspond to ROS types
base_defined_types = primitive_types | {
    # generated by gen_support()
    "MsgMetadata",
    "RosTime",
    "RosDuration"
}

def gen_metadata_item():
    """ Generate a table field containing a MsgMetadata. """ 
    yield "  __metadata:{}.MsgMetadata;".format(BASE_NS)

def gen_support():
    """ Generate supporting definitions """
    yield "// *** begin supporting definitions ***"
    # Namespace everything
    yield "namespace {};".format(BASE_NS)

    # Metadata table for all messages, to support RoboFleet
    yield "table MsgMetadata {"
    yield "  type:string;"
    yield "  topic:string;"
    yield "}"

    # All generated messages can be read as MsgWithMetadata to access metadata
    yield "table MsgWithMetadata {"
    for x in gen_metadata_item():
        yield x
    yield "}"

    # ROS time primitives
    yield "struct RosTime {"
    yield "  secs:uint32;"
    yield "  nsecs:uint32;"
    yield "}"
    yield "struct RosDuration {"
    yield "  secs:int32;"
    yield "  nsecs:int32;"
    yield "}"
    yield "// *** end supporting definitions ***"

# direct ROS to Flatbuffers type remappings
type_mapping = {
    "time": "RosTime",
    "duration": "RosDuration"
}

def type_remap(ros_type_name):
    """ apply direct translations from ROS type names to flatbuffers type names """
    if ros_type_name in type_mapping:
        return type_mapping[ros_type_name]
    return ros_type_name

class Type:
    """
    Represents a data type. Constructed from a ROS type string.
    Produces Flatbuffers type strings.

    Note on array syntax: We generate Flatbuffers *vectors*, rather than
    *arrays*. Therefore, we simply ignore any array length specified by the
    ROS definition.
    """
    def __init__(self, ros_type):
        match = re.match(r"^(?P<type>(?:(?P<ns>[\w\/]+)\/)?(?P<name>\w+))(?P<array>\[\d*\])?$", ros_type)
        if match is None:
            raise RuntimeError("Invalid ROS type: {}".format(ros_type))

        # the fully-qualified ROS type, including array syntax if any
        self.ros_type_raw = ros_type

        # the fully-qualified ROS type, without any array syntax
        self.ros_type = match.group("type")

        if match.group("ns") is not None:
            self.namespace = match.group("ns").replace("/", ".")
        else:
            self.namespace = None

        # the unqualified name of this type
        self.name = match.group("name")

        # is this an array type?
        self.is_array = match.group("array") is not None
    
    def full_namespace(self):
        """ Namespace including base namespace """
        if self.namespace is None:
            return BASE_NS
        return "{}.{}".format(BASE_NS, self.namespace)
    
    def fbs_type_name(self):
        """ Fully qualified Flatbuffers type name with remapping """
        if self.namespace is not None:
            n = "{}.{}.{}".format(BASE_NS, self.namespace, self.name)
        else: 
            n = self.name
        return type_remap(n)
    
    def fbs_type(self):
        """ Fully qualified Flatbuffers type (including array syntax) """
        full_name = self.fbs_type_name()
        if self.is_array:
            return "[{}]".format(full_name)
        return full_name

    def is_defined(self, defined_types):
        """ Determine whether this type has been defined for Flatbuffers """
        return self.fbs_type_name() in defined_types
    
    def mark_defined(self, defined_types):
        """ Add this type to defined types set """
        defined_types.add(self.fbs_type_name())

def gen_table(msg_type, items, defined_types):
    """ Generate a table for given name, type pairs. """
    yield "table {} {{".format(msg_type.name)
    for x in gen_metadata_item():
        yield x
    for k, v in items:
        attrs = ""
        # mark all non-scalar fields "required" (since they are in ROS)
        if v.fbs_type() not in scalar_types:
            attrs = " (required)"
         
        yield "  {}:{}{};".format(k, v.fbs_type(), attrs)
    yield "}"

def gen_constants_enums(msg_type):
    """
    Generate enums to represent ROS message constants.

    Enums are Flatbuffers' closest approximation of constants. We generate one
    enum per constant, with a single field "value". Only integer types are 
    supported by Flatbuffers.
    """

    from msg_util import get_msg_spec # in case the module breaks, don't break this whole script
    spec = get_msg_spec(msg_type.ros_type)

    if len(spec.constants) == 0:
        return

    yield "namespace {}.{}Constants;".format(msg_type.full_namespace(), msg_type.name)
    
    # https://google.github.io/flatbuffers/flatbuffers_guide_writing_schema.html
    allowed_enum_types = {"int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"}
    
    for c in spec.constants:
        if c.type in allowed_enum_types:
            # name is converted to lower case to avoid all-uppercase camelCase in, e.g. JS code
            yield "enum {} : {} {{ value = {} }}".format(c.name.lower(), c.type, c.val)
        else:
            print("Warning: skipped non-integral constant {}.{}".format(msg_type.fbs_type(), c.name), file=sys.stderr)

def gen_constants_table(msg_type):
    """
    Generate a table with default values to represent ROS message constants.
    Using generated Flatbuffers code, this table should be constructed as a 
    singleton to read constants values.

    The table is named MsgTypeConstants, where MsgType is the name of the given
    message type.
    
    The advantage of this method over enums is that it supports any ROS 
    constant type.
    """ 

    from msg_util import get_msg_spec # in case the module breaks, don't break this whole script
    spec = get_msg_spec(msg_type.ros_type)

    if len(spec.constants) == 0:
        return

    yield "namespace {};".format(msg_type.full_namespace())
    yield "table {}Constants {{".format(msg_type.name)
    for c in spec.constants:
        t = Type(c.type)
        if t.fbs_type() in primitive_types:
            # name is converted to lower case to avoid all-uppercase camelCase in, e.g. JS code
            yield "  {}:{} = {};".format(c.name.lower(), t.fbs_type(), c.val_text)
        else:
            raise RuntimeError("Invalid non-primitive constant")
    yield "}"

def gen_msg(msg_type, defined_types, args):
    """ 
    Generate .fbs definitions for a ROS message type, including dependencies. 
    """

    # no need to generate already-defined type
    if msg_type.is_defined(defined_types):
        raise RuntimeError("Type already generated: {}".format(msg_type.ros_type))
    msg_type.mark_defined(defined_types)

    msg_class = get_message_class(msg_type.ros_type)
    if msg_class is None:
        raise RuntimeError("ROS Message type {} not found. Ensure any necessary packages are on your path.".format(msg_type.ros_type))

    # this is officially suggested by http://wiki.ros.org/msg#Client_Library_Support
    name = msg_class._type
    keys = msg_class.__slots__
    types = [Type(name) for name in msg_class._slot_types]

    # generate dependency types
    for t in types:
        if not t.is_defined(defined_types):
            for x in gen_msg(t, defined_types, args=args):
                yield x
    
    # generate constants definitions
    if args.gen_enums:
        for x in gen_constants_enums(msg_type):
            yield x
    if args.gen_constants:
        for x in gen_constants_table(msg_type):
            yield x

    # generate type definition
    yield "namespace {};".format(msg_type.full_namespace())
    for x in gen_table(msg_type, zip(keys, types), defined_types):
        yield x

def generate_schema(msg_type_names, args):
    """ Generate Flatbuffers .fbs schema for several ROS message types, including dependencies. """
    defined_types = set(base_defined_types)
    yield "// Generated by msg2fbs; do not edit."
    for x in gen_support():
        yield x
    for msg_type_name in msg_type_names:
        for x in gen_msg(Type(msg_type_name), defined_types, args=args):
            yield x

if __name__ == "__main__":
    ap = ArgumentParser("msg2fbs.py", description=__doc__)
    ap.add_argument("message_type", nargs="+",
        help="ROS names specifying which messages to generate (e.g. std_msgs/String)")
    ap.add_argument("--output-file", "-o", nargs="?", default=None,
        help="Specify an output file. Otherwise, the schema is written to stdout.")
    ap.add_argument("--base-namespace", "-n", nargs="?", default="fb", type=str,
        help="Base namespace for Flatbuffers types, to avoid collisions with ROS types.")
    ap.add_argument("--gen-enums", "-e", action="store_true",
        help="Generate (un)signed integer-type ROS message constants as enums")
    ap.add_argument("--gen-constants", "-c", action="store_true",
        help="Generate ROS message constants as tables with default values")
    ns = ap.parse_args()

    global BASE_NS
    BASE_NS = ns.base_namespace 
    
    if ns.output_file is None:
        output_file = sys.stdout
    else:
        output_file = open(ns.output_file, "w+")
    
    lines = generate_schema(ns.message_type, args=ns)
    output_file.writelines(line + os.linesep for line in lines)
    output_file.close()
